"""
检索策略选择器
根据查询类型和特征选择最适合的检索策略
"""

from typing import Dict, List, Any
from base.logger import logger
from .query_classifier import QueryClassifier


class StrategySelector:
    """检索策略选择器"""
    
    def __init__(self):
        self.query_classifier = QueryClassifier()
        self.strategy_config = {
            'factual': {
                'top_k': 3,
                'similarity_threshold': 0.8,
                'rerank': True,
                'expand_query': False
            },
            'conceptual': {
                'top_k': 5,
                'similarity_threshold': 0.7,
                'rerank': True,
                'expand_query': True
            },
            'procedural': {
                'top_k': 7,
                'similarity_threshold': 0.6,
                'rerank': True,
                'expand_query': True
            },
            'comparative': {
                'top_k': 8,
                'similarity_threshold': 0.65,
                'rerank': True,
                'expand_query': True
            },
            'other': {
                'top_k': 5,
                'similarity_threshold': 0.7,
                'rerank': False,
                'expand_query': False
            }
        }
    
    def select_strategy(self, query: str) -> Dict[str, Any]:
        """
        根据查询选择检索策略
        
        Args:
            query: 用户查询
            
        Returns:
            检索策略配置
        """
        # 获取查询特征
        query_features = self.query_classifier.get_query_features(query)
        category = query_features['category']
        
        # 获取基础策略配置
        strategy = self.strategy_config.get(category, self.strategy_config['other']).copy()
        
        # 根据查询特征调整策略
        strategy = self._adjust_strategy_by_features(strategy, query_features)
        
        logger.info(f"选择检索策略: {category} -> {strategy}")
        return strategy
    
    def _adjust_strategy_by_features(self, strategy: Dict[str, Any], features: Dict[str, Any]) -> Dict[str, Any]:
        """
        根据查询特征调整检索策略
        
        Args:
            strategy: 基础策略配置
            features: 查询特征
            
        Returns:
            调整后的策略配置
        """
        # 根据查询长度调整
        if features['length'] > 100:  # 长查询
            strategy['top_k'] = min(strategy['top_k'] + 2, 10)
            strategy['similarity_threshold'] = max(strategy['similarity_threshold'] - 0.05, 0.5)
        elif features['length'] < 20:  # 短查询
            strategy['top_k'] = max(strategy['top_k'] - 1, 3)
            strategy['similarity_threshold'] = min(strategy['similarity_threshold'] + 0.05, 0.9)
        
        # 根据关键词数量调整
        keyword_count = len(features.get('keywords', []))
        if keyword_count > 3:
            strategy['expand_query'] = True
        elif keyword_count < 2:
            strategy['top_k'] = min(strategy['top_k'] + 1, 8)
        
        return strategy
    
    def get_retrieval_methods(self, strategy: Dict[str, Any]) -> List[str]:
        """
        根据策略获取检索方法列表
        
        Args:
            strategy: 检索策略配置
            
        Returns:
            检索方法列表
        """
        methods = ['vector_search']  # 基础向量检索
        
        if strategy.get('expand_query', False):
            methods.append('expanded_search')
        
        if strategy.get('rerank', False):
            methods.append('rerank')
        
        return methods